import os
import base64
import json
from typing import List, Optional
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from openai import OpenAI
from pydantic import BaseModel
import uvicorn

# 响应模型
class EvidenceAnalysisResponse(BaseModel):
    文件类型: str
    文件内容: str
    文件有效性说明: str
    是否可以作为证据: str

app = FastAPI(title="聊天记录证据分析API", description="分析聊天记录截图是否可以作为劳动纠纷仲裁维权证据")

# 初始化OpenAI客户端
def get_openai_client():
    api_key = os.environ.get("ARK_API_KEY")
    if not api_key:
        raise HTTPException(status_code=500, detail="ARK_API_KEY环境变量未设置")
    
    return OpenAI(
        base_url="https://ark.cn-beijing.volces.com/api/v3",
        api_key=api_key,
    )

# 将图片转换为base64编码
def image_to_base64(image_bytes: bytes) -> str:
    return base64.b64encode(image_bytes).decode('utf-8')

# 检测图片格式
def detect_image_format(image_bytes: bytes) -> str:
    """检测图片格式并返回对应的MIME类型"""
    # 检查文件头来判断图片格式
    if image_bytes.startswith(b'\xff\xd8\xff'):
        return 'image/jpeg'
    elif image_bytes.startswith(b'\x89PNG\r\n\x1a\n'):
        return 'image/png'
    elif image_bytes.startswith(b'GIF87a') or image_bytes.startswith(b'GIF89a'):
        return 'image/gif'
    elif image_bytes.startswith(b'RIFF') and b'WEBP' in image_bytes[:12]:
        return 'image/webp'
    else:
        # 默认使用jpeg
        return 'image/jpeg'

# 分析单张图片
async def analyze_single_image(client: OpenAI, image_bytes: bytes, filename: str) -> dict:
    try:
        # 检测图片格式
        image_format = detect_image_format(image_bytes)
        
        # 将图片转换为base64
        base64_image = image_to_base64(image_bytes)
        
        # 构建分析提示词
        analysis_prompt = """
请分析这张聊天记录截图，判断其是否可以作为劳动纠纷仲裁维权证据。请严格按照以下JSON格式返回结果：

{
    "文件类型": "微信聊天记录/钉钉聊天记录/飞书聊天记录/QQ聊天记录",
    "文件内容": "详细描述聊天记录的具体内容，包括参与人员、时间、关键对话内容等",
    "文件有效性说明": "详细说明该聊天记录作为证据的有效性，包括时间完整性、内容真实性、相关性等方面的分析",
    "是否可以作为证据": "是/否"
}

分析要点：
1. 识别聊天软件类型（微信、钉钉、飞书、QQ等）
2. 提取聊天内容的关键信息
3. 评估证据的完整性和真实性
4. 判断与劳动纠纷的相关性
5. 考虑法律证据的基本要求

请确保返回的是有效的JSON格式。
"""
        
        # 调用doubao模型
        response = client.chat.completions.create(
            model="doubao-seed-1-6-thinking-250615",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{image_format};base64,{base64_image}"
                            },
                        },
                        {"type": "text", "text": analysis_prompt},
                    ],
                }
            ],
        )
        
        # 解析响应
        result_text = response.choices[0].message.content
        
        # 尝试解析JSON
        try:
            result_json = json.loads(result_text)
            return result_json
        except json.JSONDecodeError:
            # 如果不是有效JSON，尝试提取JSON部分
            import re
            json_match = re.search(r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', result_text)
            if json_match:
                result_json = json.loads(json_match.group())
                return result_json
            else:
                # 如果无法解析，返回默认格式
                return {
                    "文件类型": "未识别",
                    "文件内容": result_text,
                    "文件有效性说明": "模型返回格式异常，无法正确解析",
                    "是否可以作为证据": "否"
                }
                
    except Exception as e:
        return {
            "文件类型": "分析失败",
            "文件内容": f"文件：{filename}",
            "文件有效性说明": f"分析过程中出现错误：{str(e)}",
            "是否可以作为证据": "否"
        }

@app.post("/analyze_single", response_model=EvidenceAnalysisResponse)
async def analyze_single_chat_record(file: UploadFile = File(...)):
    """
    分析单张聊天记录截图
    """
    # 验证文件类型
    if not file.content_type or not file.content_type.startswith('image/'):
        raise HTTPException(status_code=400, detail="请上传图片文件")
    
    try:
        # 读取文件内容
        image_bytes = await file.read()
        
        # 获取OpenAI客户端
        client = get_openai_client()
        
        # 分析图片
        result = await analyze_single_image(client, image_bytes, file.filename)
        
        return EvidenceAnalysisResponse(**result)
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"分析失败：{str(e)}")

@app.post("/analyze_multiple")
async def analyze_multiple_chat_records(files: List[UploadFile] = File(...)):
    """
    分析多张聊天记录截图
    """
    if len(files) == 0:
        raise HTTPException(status_code=400, detail="请至少上传一张图片")
    
    if len(files) > 10:  # 限制最多10张图片
        raise HTTPException(status_code=400, detail="最多支持同时分析10张图片")
    
    try:
        # 获取OpenAI客户端
        client = get_openai_client()
        
        results = []
        
        for file in files:
            # 验证文件类型
            if not file.content_type or not file.content_type.startswith('image/'):
                results.append({
                    "filename": file.filename,
                    "error": "文件类型不支持，请上传图片文件",
                    "analysis": None
                })
                continue
            
            try:
                # 读取文件内容
                image_bytes = await file.read()
                
                # 分析图片
                analysis_result = await analyze_single_image(client, image_bytes, file.filename)
                
                results.append({
                    "filename": file.filename,
                    "error": None,
                    "analysis": analysis_result
                })
                
            except Exception as e:
                results.append({
                    "filename": file.filename,
                    "error": f"分析失败：{str(e)}",
                    "analysis": None
                })
        
        return JSONResponse(content={
            "total_files": len(files),
            "results": results
        })
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"批量分析失败：{str(e)}")

@app.get("/")
async def root():
    return {
        "message": "聊天记录证据分析API",
        "endpoints": {
            "/analyze_single": "分析单张聊天记录截图",
            "/analyze_multiple": "分析多张聊天记录截图",
            "/docs": "API文档"
        }
    }

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8006)